import os
import sys
import torch
from diffusers import StableDiffusionPipeline
from modelscope import snapshot_download
import numpy as np
from PIL import Image
from datetime import datetime

def log(message):
    """打印带时间戳的日志"""
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print(f"[{timestamp}] {message}")

def download_model(model_id, cache_dir):
    """下载模型并返回模型路径"""
    log(f"开始下载模型 {model_id}")
    try:
        model_dir = snapshot_download(model_id, cache_dir=cache_dir)
        log(f"模型下载完成，保存在: {model_dir}")
        return model_dir
    except Exception as e:
        log(f"模型下载失败: {str(e)}")
        raise

def setup_pipeline(model_path):
    """设置和配置推理管道"""
    log("初始化生成管道...")

    # 检查是否可用GPU
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device == "cpu":
        log("警告: 使用CPU进行推理，速度会很慢")
    else:
        log(f"使用设备: {device}")

    # 加载模型
    try:
        pipe = StableDiffusionPipeline.from_pretrained(
            model_path,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            safety_checker=None,  # 禁用安全检查以提高速度
            requires_safety_checker=False
        )
        pipe = pipe.to(device)

        # 启用内存优化
        if device == "cuda":
            pipe.enable_attention_slicing()
            pipe.enable_vae_slicing()

        return pipe
    except Exception as e:
        log(f"模型加载失败: {str(e)}")
        raise

def generate_image(pipe, prompt, output_path, num_inference_steps=20):
    """生成图像"""
    log(f"使用提示词生成图像: {prompt}")
    try:
        # 生成图像
        with torch.inference_mode():
            image = pipe(
                prompt=prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=7.5,
            ).images[0]

        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # 保存图像
        image.save(output_path)
        log(f"图像已保存到: {output_path}")
        return True
    except Exception as e:
        log(f"图像生成失败: {str(e)}")
        raise

def main():
    try:
        # 设置路径
        model_base_path = 'E:/workspace/llm/text2image/models'
        model_id = 'AIkaiyuanfenxiangKK/chengxuyuan'
        output_path = os.path.join("output", "result.png")

        # 下载模型
        model_dir = download_model(model_id, model_base_path)

        # 设置pipeline
        pipe = setup_pipeline(model_dir)

        # 生成图像
        prompt = "飞流直下三千尺，油画风格，震撼壮丽的瀑布，细腻的水流，光影效果"
        generate_image(pipe, prompt, output_path, num_inference_steps=20)

        # 清理资源
        del pipe
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        log("处理完成!")
        return 0

    except Exception as e:
        log(f"发生错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())